;;;; Created on 2008-06-23 16:56:41
;;;; The environment for our agents to run around in
;;;; Trystan A Larey-Williams, 6-23-08

(in-package :neural-net)

(defun generate-maze (dimensions)
  "Generate an nth-tangular maze by recursive spatial subdivision of the given dimensions."
  (let ((maze (make-array dimensions :initial-element 0))
        (rank (length dimensions)))
    
    (labels ((make-walls (axes m n &optional (walls '()) (subdiv axes))
               (if (car axes) 
                   (let ((lst (insert (make-list (- rank 1) :initial-element #\~) (car axes) (- rank (length axes)))))
                     (space-partition-along maze lst 1 m n)
                     (make-walls (cdr axes) m n (append walls (list lst)) subdiv))          
                   (mapcar (lambda (lst)
                             (setf (apply #'aref maze 
                                          (mapcar (lambda (a c d)
                                                    (if (eql c #\~)
                                                        (interval-rand (- a 1) d)     
                                                        c))
                                                  m lst subdiv)) 0)
                             (setf (apply #'aref maze 
                                          (mapcar (lambda (b c d)
                                                    (if (eql c #\~)
                                                        (interval-rand d (+ b 1))
                                                        c))
                                                  n lst subdiv)) 0)) 
                           walls)))
                                
             (pick-axes (lst)
               (if (< (length lst) 2) 
                   nil
                   (let ((axis (interval-rand (car lst) (car (cdr lst)))))
                     (cons axis (pick-axes (cdr (cdr lst))))))))
        
      (flet ((terminate? (m n)
               (let ((terminate nil))
               (mapcar (lambda (a b) 
                         (if (< (abs (- a b)) 2.0)
                             (setf terminate t)))
                       m n)
                 terminate)))
          
        (labels ((subdivide (m n)
                   (if (not (terminate? (inc-list m 1) (dec-list n 1)))
                       (let ((subdivision (pick-axes (interleave (inc-list m 1) (dec-list n 1)))))
                         (make-walls subdivision (inc-list m 1) (dec-list n 1))
                         (mapcar #'subdivide 
                                 (combine-interchanges m subdivision)
                                 (combine-interchanges subdivision n))
                         nil))))
          (subdivide (make-list (length dimensions) :initial-element 0) 
                     (mapcar (lambda (a) (- a 1)) dimensions)))))
    
    (do ((i (- rank 1) (- i 1)))
        ((eql i -1))
      (let ((lst (insert (make-list (- rank 1) :initial-element #\~) 0 i)))
        (space-partition-along maze lst 1 (make-list rank :initial-element 0) (dec-list dimensions 1))))
    (do ((i (- rank 1) (- i 1)))
        ((eql i -1))
      (let ((lst (insert (make-list (- rank 1) :initial-element #\~) (- (nth i dimensions) 1) i)))
        (space-partition-along maze lst 1 (make-list rank :initial-element 0) (dec-list dimensions 1))))
    
    maze))


(defun render-2d-maze (maze &optional (width 640) (height 480))
  "Render a maze represented by a 2D array of binary values. Does not initialize SDL. 
Should be passed to run-sim via the pre list in order to draw properly."
  (let* ((dim (array-dimensions maze))
        (wall-width (/ width (first dim)))
        (wall-height (/ height (second dim)))
        (num-elements (* (first dim) (second dim))))
    (if (or (< wall-width 1.0) (< wall-height 1.0)) 
        (error "Surface pixel width and/or height is insufficient to render given maze."))                   
    (do ((i 0 (+ i 1)))
        ((eql i num-elements))
      (cond 
        ((eql (row-major-aref maze i) 0)
         (multiple-value-bind (x1 y1) (floor i (second dim)) 
           (sdl:draw-box-* (* x1 wall-width) (* y1 wall-height) 
                           wall-width wall-height :color sdl:*white*)))
        ((eql (row-major-aref maze i) 2)
         (multiple-value-bind (x1 y1) (floor i (second dim)) 
           (sdl:draw-box-* (* x1 wall-width) (* y1 wall-height) 
                           wall-width wall-height :color sdl:*red*)))))))
             

(defun run-sim (x y pre loop post)
  "Sets up a x by y SDL window and simulation loop. Runs all functions in the pre prior
to entering the game loop (but after sdl init), all loop functions in each game loop iteration and 
all post functions after a quit event is received."
  (sdl:with-init (sdl:sdl-init-video) 
                 (sdl:window x y :title-caption "Agent Simulation")
                 (setf (sdl:frame-rate) 60)    
                 (flet ((do-funs (lst)
                          (mapcar (lambda (e)
                                    (if (typep e 'list)
                                        (apply (car e) (cdr e))
                                        (funcall e)))
                                  lst)))          
                 (do-funs pre)
                 (sdl:with-events ()
                                  (:QUIT-EVENT () 
                                    (do-funs post)
                                    T)
                                  (:IDLE ()
                                    (do-funs loop)
                                    (sdl:update-display))))))


(defun lab-rat (maze position cheese)
  "The majority of the code here constitures the input and output mapping functions. 
That is, the functions that map the input and output layers of the NN to the environment.
The last few lines or so actually build the NN for our agent."
  (let ((old-dist 1000000000))
    (flet ((input-map (input-layer)
             (let ((posx (nth 0 position))
                   (posy (nth 1 position))
                   (layer (neurons input-layer))
                   (new-dist (lisp-utils:distance position cheese)))
               (setf (activation-state (nth 0 layer)) (aref maze (- posx 1) posy))
               (setf (activation-state (nth 1 layer)) (aref maze (- posx 1) (+ posy 1)))
               (setf (activation-state (nth 2 layer)) (aref maze posx (+ posy 1)))
               (setf (activation-state (nth 3 layer)) (aref maze (+ posx 1) (+ posy 1)))
               (setf (activation-state (nth 4 layer)) (aref maze (+ posx 1) posy))
               (setf (activation-state (nth 5 layer)) (aref maze (+ posx 1) (- posy 1)))
               (setf (activation-state (nth 6 layer)) (aref maze posx (- posy 1)))
               (setf (activation-state (nth 7 layer)) (aref maze (- posx 1) (- posy 1)))
               (if (< new-dist old-dist)
                   (setf (activation-state (nth 8 layer)) 1)
                   (setf (activation-state (nth 8 layer)) 0))
               (setf old-dist new-dist)))
           (output-map (output-layer)
             ; sum the output layers (sum => intended direction 0-7) 
             (let ((output (reduce #'+ (mapcar (lambda (n)
                                   (previous-activation n))
                                 (neurons output-layer))))
                   (posx (nth 0 position))
                   (posy (nth 1 position)))
               (cond 
                 ((and (< output 1) (/= (aref maze (- posx 1) posy) 1)) (setf position (list (- posx 1) posy)))
                 ((and (< output 2) (/= (aref maze (- posx 1) (+ posy 1)) 1)) (setf position (list (- posx 1) (+ posy 1))))
                 ((and (< output 3) (/= (aref maze posx (+ posy 1)) 1)) (setf position (list posx (+ posy 1))))
                 ((and (< output 4) (/= (aref maze (+ posx 1) (+ posy 1)) 1)) (setf position (list (+ posx 1) (+ posy 1))))
                 ((and (< output 5) (/= (aref maze (+ posx 1) posy) 1)) (setf position (list (+ posx 1) posy)))
                 ((and (< output 6) (/= (aref maze (+ posx 1) (- posy 1)) 1)) (setf position (list (+ posx 1) (- posy 1))))
                 ((and (< output 7) (/= (aref maze posx (- posy 1)) 1)) (setf position (list posx (- posy 1))))
                 ((and (< output 8) (/= (aref maze (- posx 1) (- posy 1)) 1)) (setf position (list (- posx 1) (- posy 1)))))
               (print output)
               (setf (apply #'aref maze position) 2))))
      (let ((l-in (make-neural-layer 9 #'make-neuron )) 
            (l-feedback (make-neural-layer 8 #'make-neuron ))
            (l-out (make-neural-layer 8 #'make-neuron )))  
        (wire-feedforward (list l-in l-out l-feedback l-out))
        (make-feedforward-nn l-in l-out (list (list l-in l-feedback) l-out) #'input-map #'output-map )))))


(defun rat-in-maze ( maze-dim )
  (let* ((x (nth 0 maze-dim))
         (y (nth 1 maze-dim))
         (maze (generate-maze maze-dim))
         (rat (lab-rat maze '(1 1) (list (- x 1) (- y 1)))))
    (run-sim 640
             480
             nil 
             (list (list #'run-nn-pass rat) (list #'render-2d-maze maze)) 
             nil)))
